{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021-2022 @ Shenzhen Bay Laboratory & Peking University & Huawei Technologies Co., Ltd\n",
    "\n",
    "This code is a part of Cybertron package.\n",
    "\n",
    "The Cybertron is open-source software based on the AI-framework:\n",
    "MindSpore (https://www.mindspore.cn/)\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "\n",
    "You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "Cybertron tutorial 00: The preprocessing of dataset\n",
    "\n",
    "The NPZ file can be download from:\n",
    "dataset_qm9.npz: http://gofile.me/6Utp7/tJ5hoDIAo\n",
    "ethanol_dft.npz: http://gofile.me/6Utp7/hbQBofAFM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(969479:140204621940544,MainProcess):2022-08-10-16:18:04.851.759 [mindspore/run_check/_check_version.py:137] Can not found cuda libs, please confirm that the correct cuda version has been installed, you can refer to the installation guidelines: https://www.mindspore.cn/install\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from cybertron.dataset import DatasetProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atom_ref (10, 6)\n",
      "valid_idx (100,)\n",
      "uncharacterized (3054,)\n",
      "excluded (11,)\n",
      "Nos (133885, 1)\n",
      "Natom (133885, 1)\n",
      "SMILES (133885, 2)\n",
      "InChI (133885, 2)\n",
      "Z (133885, 29)\n",
      "R (133885, 29, 3)\n",
      "Q (133885, 29)\n",
      "properties (133885, 15)\n",
      "frequencies (133885, 150)\n"
     ]
    }
   ],
   "source": [
    "qm9_file = 'dataset_qm9.npz'\n",
    "qm9_data = np.load(qm9_file)\n",
    "for k, v in qm9_data.items():\n",
    "    print(k, v.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========================================================================================================\n",
      "Number of atoms: 29\n",
      "Number of data: 133885\n",
      "Dimension of space: 3\n",
      "Total number of effective atoms: 2407753\n",
      "Shape of atom types (Z): (133885, 29)\n",
      "Shape of position (R): (133885, 29, 3)\n",
      "Shape of label (E): (133885, 12)\n",
      "Reference value for atom types:\n",
      "   Type     Label0    Label1    Label2    Label3    Label4    Label5    Label6    Label7    Label8    Label9   Label10   Label11\n",
      "   0:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   1:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00 -5.00e-01 -4.99e-01 -4.98e-01 -5.11e-01  2.98e+00\n",
      "   2:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   3:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   4:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   5:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00\n",
      "   6:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00 -3.78e+01 -3.78e+01 -3.78e+01 -3.79e+01  2.98e+00\n",
      "   7:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00 -5.46e+01 -5.46e+01 -5.46e+01 -5.46e+01  2.98e+00\n",
      "   8:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00 -7.51e+01 -7.51e+01 -7.51e+01 -7.51e+01  2.98e+00\n",
      "   9:        0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00  0.00e+00 -9.97e+01 -9.97e+01 -9.97e+01 -9.97e+01  2.98e+00\n",
      "Length unit: Angstrom\n",
      "Energy unit: Hartree\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to graph:\n",
      "   ID              Mean             STD             Min             Max          Median           Range\n",
      "    0:     2.706037e+00    1.530388e+00    0.000000e+00    2.955640e+01    2.500000e+00    2.955640e+01\n",
      "    1:     7.519130e+01    8.187762e+00    6.310000e+00    1.966200e+02    7.550000e+01    1.903100e+02\n",
      "    2:    -2.399767e-01    2.213135e-02   -4.286000e-01   -1.017000e-01   -2.410000e-01    3.269000e-01\n",
      "    3:     1.112377e-02    4.693589e-02   -1.750000e-01    1.935000e-01    1.200000e-02    3.685000e-01\n",
      "    4:     2.511004e-01    4.751871e-02    2.460000e-02    6.221000e-01    2.494000e-01    5.975000e-01\n",
      "    5:     1.189527e+03    2.797561e+02    1.900020e+01    3.374753e+03    1.147586e+03    3.355753e+03\n",
      "    6:     1.485244e-01    3.327377e-02    1.595100e-02    2.739440e-01    1.483290e-01    2.579930e-01\n",
      "    7:    -2.790100e+00    3.813690e-01   -4.156831e+00   -3.395770e-01   -2.794326e+00    3.817254e+00\n",
      "    8:    -2.807100e+00    3.847518e-01   -4.185451e+00   -3.409900e-01   -2.811259e+00    3.844461e+00\n",
      "    9:    -2.823142e+00    3.874847e-01   -4.211903e+00   -3.428790e-01   -2.827292e+00    3.869024e+00\n",
      "   10:    -2.596597e+00    3.509213e-01   -3.851932e+00   -3.209630e-01   -2.601118e+00    3.530969e+00\n",
      "   11:    -2.200885e+01    6.126555e+00   -4.304800e+01    6.808000e+00   -2.214800e+01    4.985600e+01\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to each atom:\n",
      "   ID              Mean             STD             Min             Max          Median           Range\n",
      "    0:     1.599993e-01    1.036793e-01    0.000000e+00    1.617757e+00    1.398750e-01    1.617757e+00\n",
      "    1:     4.241285e+00    5.026749e-01    2.103333e+00    1.194200e+01    4.140000e+00    9.838667e+00\n",
      "    2:    -1.378340e-02    3.215979e-03   -1.201333e-01   -4.750000e-03   -1.314211e-02    1.153833e-01\n",
      "    3:     3.423312e-04    2.804465e-03   -2.603333e-02    3.870000e-02    6.764706e-04    6.473333e-02\n",
      "    4:     1.412572e-02    2.672984e-03    1.118182e-03    1.265333e-01    1.397778e-02    1.254152e-01\n",
      "    5:     6.691506e+01    1.522971e+01    6.333400e+00    2.228472e+02    6.424726e+01    2.165138e+02\n",
      "    6:     8.172795e-03    5.800831e-04    3.340556e-03    9.492037e-03    8.260947e-03    6.151481e-03\n",
      "    7:    -1.559705e-01    6.967993e-03   -2.138314e-01   -1.103017e-01   -1.550146e-01    1.035296e-01\n",
      "    8:    -1.569094e-01    6.941453e-03   -2.143701e-01   -1.110025e-01   -1.559679e-01    1.033676e-01\n",
      "    9:    -1.577998e-01    6.934442e-03   -2.151961e-01   -1.117112e-01   -1.568604e-01    1.034849e-01\n",
      "   10:    -1.452039e-01    6.901672e-03   -2.021993e-01   -1.033207e-01   -1.442428e-01    9.887850e-02\n",
      "   11:    -1.203925e+00    1.922459e-01   -1.782231e+00    7.564444e-01   -1.225316e+00    2.538675e+00\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of data: 133885\n",
      "=========================================================================================================\n"
     ]
    }
   ],
   "source": [
    "atom_ref = qm9_data['atom_ref']\n",
    "ref = np.zeros((atom_ref.shape[0], 6))\n",
    "atom_ref = np.concatenate((ref, atom_ref), axis=-1)\n",
    "\n",
    "ds_qm9 = DatasetProcessor(\n",
    "    name='qm9',\n",
    "    atom_types=qm9_data['Z'],\n",
    "    position=qm9_data['R'],\n",
    "    label=qm9_data['properties'][:, 3:],\n",
    "    length_unit='A',\n",
    "    energy_unit='Ha',\n",
    "    type_ref=atom_ref,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Remove 3062 samples from 130823 dataset.\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to graph:\n",
      "   ID              Mean             STD             Min             Max          Median           Range\n",
      "    0:     2.672827e+00    1.503183e+00    0.000000e+00    2.955640e+01    2.475200e+00    2.955640e+01\n",
      "    1:     7.528103e+01    8.173763e+00    6.310000e+00    1.966200e+02    7.560000e+01    1.903100e+02\n",
      "    2:    -2.402090e-01    2.196555e-02   -4.286000e-01   -1.017000e-01   -2.411000e-01    3.269000e-01\n",
      "    3:     1.184083e-02    4.684326e-02   -1.750000e-01    1.935000e-01    1.260000e-02    3.685000e-01\n",
      "    4:     2.520497e-01    4.718871e-02    2.460000e-02    6.221000e-01    2.502000e-01    5.975000e-01\n",
      "    5:     1.189402e+03    2.804632e+02    1.900020e+01    3.374753e+03    1.147221e+03    3.355753e+03\n",
      "    6:     1.490930e-01    3.313445e-02    1.595100e-02    2.739440e-01    1.486310e-01    2.579930e-01\n",
      "    7:    -2.797248e+00    3.793542e-01   -4.156831e+00   -3.395770e-01   -2.802355e+00    3.817254e+00\n",
      "    8:    -2.814318e+00    3.827137e-01   -4.185451e+00   -3.409900e-01   -2.819457e+00    3.844461e+00\n",
      "    9:    -2.830406e+00    3.854364e-01   -4.211903e+00   -3.428790e-01   -2.835493e+00    3.869024e+00\n",
      "   10:    -2.603233e+00    3.490228e-01   -3.851932e+00   -3.209630e-01   -2.608446e+00    3.530969e+00\n",
      "   11:    -2.213512e+01    6.082039e+00   -4.304800e+01    9.420000e-01   -2.226200e+01    4.399000e+01\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to each atom:\n",
      "   ID              Mean             STD             Min             Max          Median           Range\n",
      "    0:     1.576426e-01    1.020934e-01    0.000000e+00    1.617757e+00    1.378867e-01    1.617757e+00\n",
      "    1:     4.233573e+00    4.964381e-01    2.103333e+00    1.194200e+01    4.134667e+00    9.838667e+00\n",
      "    2:    -1.375555e-02    3.198124e-03   -1.201333e-01   -4.750000e-03   -1.311875e-02    1.153833e-01\n",
      "    3:     3.861658e-04    2.786526e-03   -2.603333e-02    3.870000e-02    7.058824e-04    6.473333e-02\n",
      "    4:     1.414170e-02    2.665319e-03    1.118182e-03    1.265333e-01    1.399375e-02    1.254152e-01\n",
      "    5:     6.669718e+01    1.513921e+01    6.333400e+00    2.228472e+02    6.403401e+01    2.165138e+02\n",
      "    6:     8.183581e-03    5.721176e-04    3.429400e-03    9.492037e-03    8.269556e-03    6.062637e-03\n",
      "    7:    -1.559400e-01    6.950499e-03   -2.138314e-01   -1.103017e-01   -1.549875e-01    1.035296e-01\n",
      "    8:    -1.568803e-01    6.924694e-03   -2.143701e-01   -1.110025e-01   -1.559435e-01    1.033676e-01\n",
      "    9:    -1.577709e-01    6.917738e-03   -2.151961e-01   -1.117112e-01   -1.568374e-01    1.034849e-01\n",
      "   10:    -1.451742e-01    6.883536e-03   -2.021993e-01   -1.033207e-01   -1.442178e-01    9.887850e-02\n",
      "   11:    -1.208249e+00    1.889129e-01   -1.782231e+00    1.046667e-01   -1.229000e+00    1.886897e+00\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of data: 130823\n",
      "=========================================================================================================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([     0,      1,      2, ..., 133882, 133883, 133884], dtype=int32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "exc_idx = np.concatenate(\n",
    "    (qm9_data['excluded'], qm9_data['uncharacterized']), -1) - 1\n",
    "ds_qm9.exclude_data(exc_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Change the length unit from \"Angstrom\" to \"nm\"\n",
      "Change the energy unit from \"Hartree\" to \"kJ mol-1\"\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to graph:\n",
      "   ID              Mean             STD             Min             Max          Median           Range\n",
      "    0:     2.672827e+00    1.503183e+00    0.000000e+00    2.955640e+01    2.475200e+00    2.955640e+01\n",
      "    1:     7.528103e+01    8.173763e+00    6.310000e+00    1.966200e+02    7.560000e+01    1.903100e+02\n",
      "    2:    -6.306687e+02    5.767056e+01   -1.125289e+03   -2.670134e+02   -6.330081e+02    8.582760e+02\n",
      "    3:     3.108809e+01    1.229870e+02   -4.594625e+02    5.080343e+02    3.308130e+01    9.674968e+02\n",
      "    4:     6.617564e+02    1.238940e+02    6.458730e+01    1.633324e+03    6.569002e+02    1.568736e+03\n",
      "    5:     1.189402e+03    2.804632e+02    1.900020e+01    3.374753e+03    1.147221e+03    3.355753e+03\n",
      "    6:     3.914438e+02    8.699451e+01    4.187935e+01    7.192400e+02    3.902307e+02    6.773607e+02\n",
      "    7:    -7.344176e+03    9.959946e+02   -1.091376e+04   -8.915595e+02   -7.357584e+03    1.002220e+04\n",
      "    8:    -7.388992e+03    1.004815e+03   -1.098890e+04   -8.952693e+02   -7.402485e+03    1.009363e+04\n",
      "    9:    -7.431231e+03    1.011963e+03   -1.105835e+04   -9.002289e+02   -7.444587e+03    1.015812e+04\n",
      "   10:    -6.834789e+03    9.163595e+02   -1.011325e+04   -8.426884e+02   -6.848475e+03    9.270560e+03\n",
      "   11:    -2.213512e+01    6.082039e+00   -4.304800e+01    9.420000e-01   -2.226200e+01    4.399000e+01\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to each atom:\n",
      "   ID              Mean             STD             Min             Max          Median           Range\n",
      "    0:     1.576426e-01    1.020934e-01    0.000000e+00    1.617757e+00    1.378867e-01    1.617757e+00\n",
      "    1:     4.233573e+00    4.964381e-01    2.103333e+00    1.194200e+01    4.134667e+00    9.838667e+00\n",
      "    2:    -3.611519e+01    8.396674e+00   -3.154101e+02   -1.247113e+01   -3.444328e+01    3.029390e+02\n",
      "    3:     1.013878e+00    7.316026e+00   -6.835052e+01    1.016069e+02    1.853294e+00    1.699574e+02\n",
      "    4:     3.712904e+01    6.997796e+00    2.935787e+00    3.322133e+02    3.674059e+01    3.292775e+02\n",
      "    5:     6.669718e+01    1.513921e+01    6.333400e+00    2.228472e+02    6.403401e+01    2.165138e+02\n",
      "    6:     2.148599e+01    1.502095e+00    9.003890e+00    2.492135e+01    2.171172e+01    1.591745e+01\n",
      "    7:    -4.094204e+02    1.824854e+01   -5.614143e+02   -2.895973e+02   -4.069196e+02    2.718171e+02\n",
      "    8:    -4.118893e+02    1.818079e+01   -5.628288e+02   -2.914371e+02   -4.094298e+02    2.713917e+02\n",
      "    9:    -4.142276e+02    1.816252e+01   -5.649975e+02   -2.932979e+02   -4.117766e+02    2.716996e+02\n",
      "   10:    -3.811549e+02    1.807273e+01   -5.308742e+02   -2.712686e+02   -3.786439e+02    2.596055e+02\n",
      "   11:    -1.208249e+00    1.889129e-01   -1.782231e+00    1.046667e-01   -1.229000e+00    1.886897e+00\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of data: 130823\n",
      "=========================================================================================================\n"
     ]
    }
   ],
   "source": [
    "ds_qm9.convert_units('nm', 'kj/mol', [2, 3, 4, 6, 7, 8, 9, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([106218,  89897, 119180, ...,  36446,   4509,  42720], dtype=int32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_qm9.shuffle_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   ID.            Scale           Shift        Mode\n",
      "    0.     1.503183e+00    2.672827e+00       graph\n",
      "    1.     8.173763e+00    7.528103e+01       graph\n",
      "    2.     5.767056e+01   -6.306687e+02       graph\n",
      "    3.     1.229870e+02    3.108809e+01       graph\n",
      "    4.     1.238940e+02    6.617564e+02       graph\n",
      "    5.     2.804632e+02    1.189402e+03       graph\n",
      "    6.     8.699451e+01    3.914438e+02       graph\n",
      "    7.     1.824854e+01   -4.094204e+02    atomwise\n",
      "    8.     1.818079e+01   -4.118893e+02    atomwise\n",
      "    9.     1.816252e+01   -4.142276e+02    atomwise\n",
      "   10.     1.807273e+01   -3.811549e+02    atomwise\n",
      "   11.     6.082039e+00   -2.213512e+01       graph\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of Training, Validation and Test dataset: 1024, 128, 1024.\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Dataset file:\n",
      "   Training dataset    (normalized):   dataset_qm9_normed_trainset_1024.npz\n",
      "   Validation dataset  (normalized):   dataset_qm9_normed_validset_128.npz\n",
      "   Test dataset        (normalized):   dataset_qm9_normed_testset_1024.npz\n",
      "=========================================================================================================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'name': 'qm9',\n",
       "  'num_atoms': 29,\n",
       "  'num_dataset': 1024,\n",
       "  'data_index': array([108284,  91708, 121372, ...,  16621,  36552,  50673], dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'type_ref': array([[ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.3134669e+03,\n",
       "          -1.3097491e+03, -1.3072681e+03, -1.3414390e+03,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -9.9366711e+04,\n",
       "          -9.9362984e+04, -9.9360508e+04, -9.9404898e+04,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.4330994e+05,\n",
       "          -1.4330622e+05, -1.4330373e+05, -1.4334942e+05,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.9708206e+05,\n",
       "          -1.9707834e+05, -1.9707588e+05, -1.9712133e+05,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -2.6181155e+05,\n",
       "          -2.6180783e+05, -2.6180534e+05, -2.6185044e+05,  2.9809999e+00]],\n",
       "        dtype=float32),\n",
       "  'Z': array([[6, 6, 6, ..., 0, 0, 0],\n",
       "         [8, 6, 6, ..., 0, 0, 0],\n",
       "         [8, 6, 6, ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [8, 6, 6, ..., 0, 0, 0]], dtype=int32),\n",
       "  'R': array([[[-1.1142875e-02,  1.5812907e-01,  4.4413228e-03],\n",
       "          [ 3.5342858e-03,  5.6812940e-03, -4.7988514e-03],\n",
       "          [ 7.2577603e-02, -5.4099612e-02,  1.1570221e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 6.1214771e-03,  1.4215238e-01, -1.3827609e-02],\n",
       "          [ 1.5041523e-03,  2.0087317e-03,  1.4307916e-03],\n",
       "          [ 6.6030197e-02, -5.1852573e-02,  1.3371411e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 2.0818852e-03,  1.4365010e-01, -2.3040319e-02],\n",
       "          [ 5.3538070e-03,  2.5866476e-03, -6.2697297e-03],\n",
       "          [-1.3854927e-01, -4.7261249e-02, -2.3530282e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[-1.0684761e-02,  1.5484959e-01, -1.2202791e-03],\n",
       "          [ 1.9219632e-03,  1.7879457e-03,  7.7857250e-05],\n",
       "          [ 7.9916753e-02, -5.3879716e-02, -1.1918602e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 1.5979379e-02, -1.2394467e-01, -3.5287064e-02],\n",
       "          [ 1.2903147e-02, -3.8187262e-03, -3.7626967e-02],\n",
       "          [ 1.2225902e-02,  1.4184269e-01, -3.7501145e-02],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 2.0368034e-02, -2.0850457e-02, -1.1839227e-02],\n",
       "          [ 1.0124510e-02,  9.8646425e-02, -1.7432189e-03],\n",
       "          [ 1.2764639e-01,  1.9380136e-01,  5.1811035e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]]], dtype=float32),\n",
       "  'E': array([[ 1.4071294e+00, -3.5858968e-01,  1.2788050e-01, ...,\n",
       "           3.2806509e+00,  3.4836495e+00, -4.4200704e-03],\n",
       "         [ 4.7244608e-01, -6.0205162e-01, -1.1340642e-01, ...,\n",
       "           1.5441908e-03,  6.3511372e-01, -4.6676499e-01],\n",
       "         [-4.9569961e-01,  7.8042066e-01, -7.9629391e-01, ...,\n",
       "           3.0543846e+01,  3.0450903e+01, -1.5323615e+00],\n",
       "         ...,\n",
       "         [-8.7975138e-01,  1.9562265e-01,  8.2442576e-01, ...,\n",
       "           8.3192816e+00,  8.2801971e+00, -7.3690468e-01],\n",
       "         [-6.5176845e-01,  1.4031448e+00, -5.0037605e-01, ...,\n",
       "          -1.5616557e+01, -1.6529350e+01,  8.7341052e-01],\n",
       "         [ 9.7777385e-01,  7.3279977e-02, -8.6913526e-01, ...,\n",
       "          -1.4840591e+01, -1.5065827e+01, -6.5090515e-02]], dtype=float32),\n",
       "  'scale': array([  1.5031829,   8.173762 ,  57.670555 , 122.986984 , 123.89397  ,\n",
       "         280.4632   ,  86.99451  ,  18.248537 ,  18.180786 ,  18.162521 ,\n",
       "          18.072725 ,   6.0820394], dtype=float32),\n",
       "  'shift': array([   2.6728272,   75.28103  , -630.6687   ,   31.08809  ,\n",
       "          661.7564   , 1189.402    ,  391.44376  , -409.42038  ,\n",
       "         -411.88934  , -414.2276   , -381.15494  ,  -22.135117 ],\n",
       "        dtype=float32),\n",
       "  'scaleshift_mode': ('graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'graph'),\n",
       "  'use_atomwise_scaleshift': array([False, False, False, False, False, False, False,  True,  True,\n",
       "          True,  True, False])},\n",
       " {'name': 'qm9',\n",
       "  'num_atoms': 29,\n",
       "  'num_dataset': 128,\n",
       "  'data_index': array([119217,  83638,  17832,  73544,  72779, 133356,  52696,   4530,\n",
       "          48210,  66290,  88326, 126158,  64508,  22979,  85592,  52201,\n",
       "          38192,  13523,   5426,  41008, 121083,  37617,  58069, 127127,\n",
       "          92364,  65602, 133548,  50618,  91513, 107505,  12987,  39937,\n",
       "         123078,  44214,  39705,  95804,  70470, 100750,  33904, 116382,\n",
       "         109327,  51323,  71322,   2226,  68799,  44729,  52858,  94200,\n",
       "          51818,  10788, 102755,  62127,  36051,  22192,  26037,  50578,\n",
       "          37878, 117692, 128516,  14380,  56790, 108351, 116105,  75064,\n",
       "          43090, 115079,  77340,  15407, 113247, 105654,  88470, 112237,\n",
       "          39623,  34637, 132744,  16337, 130728,  76492, 115342,  83963,\n",
       "         113172, 109162, 111437,  12212, 122438, 130696, 126121,  12711,\n",
       "           7154,  64681,  50415,  64450,   8653, 132978,  14629,  51548,\n",
       "         105042,   2330,  20441,  62091,  40407,  97077,  47239,  41513,\n",
       "          95872, 103687,  80413,   2061,  32005, 108198,  97737, 113544,\n",
       "           3643,  81688, 117340,  11574,  47601, 125181, 119468, 110975,\n",
       "          53608,  32541,  59988,   6177,  44991, 109206, 128630,  21296],\n",
       "        dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'type_ref': array([[ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.3134669e+03,\n",
       "          -1.3097491e+03, -1.3072681e+03, -1.3414390e+03,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -9.9366711e+04,\n",
       "          -9.9362984e+04, -9.9360508e+04, -9.9404898e+04,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.4330994e+05,\n",
       "          -1.4330622e+05, -1.4330373e+05, -1.4334942e+05,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.9708206e+05,\n",
       "          -1.9707834e+05, -1.9707588e+05, -1.9712133e+05,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -2.6181155e+05,\n",
       "          -2.6180783e+05, -2.6180534e+05, -2.6185044e+05,  2.9809999e+00]],\n",
       "        dtype=float32),\n",
       "  'Z': array([[6, 8, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 8, ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [8, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 8, 6, ..., 0, 0, 0],\n",
       "         [6, 8, 6, ..., 0, 0, 0]], dtype=int32),\n",
       "  'R': array([[[-2.5081968e-02,  1.3415512e-01,  2.8292919e-04],\n",
       "          [-2.2156516e-02, -6.0766470e-03, -1.3487363e-02],\n",
       "          [ 7.7686816e-02, -5.0499283e-02, -1.0191814e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-6.0903680e-02,  1.3782467e-01,  9.7752549e-03],\n",
       "          [-5.2218386e-03, -3.0788200e-03, -4.9906294e-03],\n",
       "          [ 7.6597616e-02, -4.1184250e-02, -1.3205290e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-7.8710569e-03,  1.4489976e-01,  3.4626110e-03],\n",
       "          [-2.3963023e-03, -6.7769708e-03, -5.4659281e-04],\n",
       "          [ 9.2756651e-02, -7.0822537e-02,  3.8601402e-02],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[ 9.3458230e-03,  1.3364875e-01,  1.6986379e-02],\n",
       "          [ 1.0132672e-03, -4.4945139e-03, -4.9090753e-03],\n",
       "          [-1.4174549e-01, -5.1413324e-02, -8.1585599e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-3.5088016e-03,  1.3995786e-01,  4.4502923e-03],\n",
       "          [ 2.8237996e-03, -3.7284407e-03, -4.7976729e-03],\n",
       "          [ 6.2962823e-02, -5.3789377e-02, -1.1000546e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-3.2752431e-03,  1.4069633e-01,  2.2779095e-03],\n",
       "          [-1.2417695e-03, -8.2001981e-04, -1.8107735e-02],\n",
       "          [-1.1851490e-01, -6.3027807e-02, -3.6884123e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]]], dtype=float32),\n",
       "  'E': array([[ 7.96292245e-01,  3.85253787e-01, -1.23334193e+00, ...,\n",
       "           1.39347010e+01,  1.36809101e+01,  1.23333134e-01],\n",
       "         [-1.48985672e+00,  2.61687696e-01,  3.73720020e-01, ...,\n",
       "           2.01012917e+01,  2.03667793e+01, -8.54792774e-01],\n",
       "         [ 4.66857940e-01, -5.70242524e-01,  3.28194171e-01, ...,\n",
       "           5.12374973e+00,  4.78643751e+00, -8.02862400e-04],\n",
       "         ...,\n",
       "         [ 5.76358914e-01, -6.26520157e-01, -9.06435028e-02, ...,\n",
       "           3.27357817e+00,  3.31974936e+00,  4.11887676e-01],\n",
       "         [ 1.12772238e+00, -1.84138286e+00, -2.06191206e+00, ...,\n",
       "          -8.11202717e+00, -7.96791697e+00,  2.12989688e+00],\n",
       "         [-6.00543797e-01, -1.27860653e+00, -1.36169344e-01, ...,\n",
       "          -1.34799826e+00, -1.71013987e+00,  8.05670083e-01]], dtype=float32),\n",
       "  'scale': array([  1.5031829,   8.173762 ,  57.670555 , 122.986984 , 123.89397  ,\n",
       "         280.4632   ,  86.99451  ,  18.248537 ,  18.180786 ,  18.162521 ,\n",
       "          18.072725 ,   6.0820394], dtype=float32),\n",
       "  'shift': array([   2.6728272,   75.28103  , -630.6687   ,   31.08809  ,\n",
       "          661.7564   , 1189.402    ,  391.44376  , -409.42038  ,\n",
       "         -411.88934  , -414.2276   , -381.15494  ,  -22.135117 ],\n",
       "        dtype=float32),\n",
       "  'scaleshift_mode': ('graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'graph'),\n",
       "  'use_atomwise_scaleshift': array([False, False, False, False, False, False, False,  True,  True,\n",
       "          True,  True, False])},\n",
       " {'name': 'qm9',\n",
       "  'num_atoms': 29,\n",
       "  'num_dataset': 1024,\n",
       "  'data_index': array([  9948, 110902,  41365, ...,  18549,  78512,  74618], dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'type_ref': array([[ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.3134669e+03,\n",
       "          -1.3097491e+03, -1.3072681e+03, -1.3414390e+03,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -9.9366711e+04,\n",
       "          -9.9362984e+04, -9.9360508e+04, -9.9404898e+04,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.4330994e+05,\n",
       "          -1.4330622e+05, -1.4330373e+05, -1.4334942e+05,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -1.9708206e+05,\n",
       "          -1.9707834e+05, -1.9707588e+05, -1.9712133e+05,  2.9809999e+00],\n",
       "         [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00,  0.0000000e+00,  0.0000000e+00, -2.6181155e+05,\n",
       "          -2.6180783e+05, -2.6180534e+05, -2.6185044e+05,  2.9809999e+00]],\n",
       "        dtype=float32),\n",
       "  'Z': array([[6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [8, 6, 7, ..., 0, 0, 0],\n",
       "         [6, 6, 8, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0]], dtype=int32),\n",
       "  'R': array([[[-1.62902959e-02,  1.36730880e-01, -9.98675730e-03],\n",
       "          [ 1.38593940e-02, -1.06998775e-02,  4.15915297e-03],\n",
       "          [ 1.52657643e-01, -6.42142296e-02, -8.41158908e-03],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[ 1.10524502e-02,  1.58110440e-01,  1.08578904e-02],\n",
       "          [-4.21275035e-04,  5.68928989e-03,  4.74962144e-04],\n",
       "          [ 6.19037785e-02, -5.18585853e-02, -1.26719937e-01],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[-8.54106329e-04,  1.45507395e-01, -3.53108719e-02],\n",
       "          [ 1.44421622e-01,  1.60085946e-01,  2.47270558e-02],\n",
       "          [ 1.42592028e-01,  8.34128074e-03,  5.83807863e-02],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[-9.12077143e-04, -1.42512412e-03, -8.75520986e-03],\n",
       "          [ 5.50002500e-04,  1.18384898e-01, -2.63127428e-03],\n",
       "          [-9.84417796e-02,  2.13059574e-01,  8.81935563e-03],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[ 6.42315764e-03,  1.41510636e-01, -1.73014812e-02],\n",
       "          [ 3.31989228e-04, -6.93340274e-03, -2.68555875e-03],\n",
       "          [ 5.51430434e-02, -9.14408639e-02, -1.08960524e-01],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[-2.74167373e-03,  1.50980234e-01,  2.23346078e-03],\n",
       "          [-1.16161034e-02,  1.88250549e-03, -3.42915463e-03],\n",
       "          [-1.13114680e-03, -7.84160644e-02, -1.10561773e-01],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]]],\n",
       "        dtype=float32),\n",
       "  'E': array([[ -0.01884477,  -0.7256177 ,  -0.32737786, ...,  11.401514  ,\n",
       "           10.792538  ,   0.2995898 ],\n",
       "         [ -0.51905006,   0.24945343,  -0.15893225, ...,  18.104107  ,\n",
       "           18.596178  ,  -1.2431494 ],\n",
       "         [ -0.7543508 ,  -0.22034244,  -0.04056509, ...,  -6.2811456 ,\n",
       "           -5.6343384 ,  -0.57001984],\n",
       "         ...,\n",
       "         [  0.69949764,  -1.0950925 ,   0.27811575, ..., -18.439135  ,\n",
       "          -18.475504  ,   0.35072398],\n",
       "         [ -0.2486904 ,   0.05248172,   0.79255766, ...,  10.689566  ,\n",
       "           10.988834  ,  -0.11655352],\n",
       "         [ -0.8993099 ,   0.5712147 ,   1.0975808 , ...,  -8.230782  ,\n",
       "           -7.856421  ,  -0.4499943 ]], dtype=float32),\n",
       "  'scale': array([  1.5031829,   8.173762 ,  57.670555 , 122.986984 , 123.89397  ,\n",
       "         280.4632   ,  86.99451  ,  18.248537 ,  18.180786 ,  18.162521 ,\n",
       "          18.072725 ,   6.0820394], dtype=float32),\n",
       "  'shift': array([   2.6728272,   75.28103  , -630.6687   ,   31.08809  ,\n",
       "          661.7564   , 1189.402    ,  391.44376  , -409.42038  ,\n",
       "         -411.88934  , -414.2276   , -381.15494  ,  -22.135117 ],\n",
       "        dtype=float32),\n",
       "  'scaleshift_mode': ('graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'graph'),\n",
       "  'use_atomwise_scaleshift': array([False, False, False, False, False, False, False,  True,  True,\n",
       "          True,  True, False])})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mode = ['graph', 'graph', 'graph', 'graph', 'graph', 'graph',\n",
    "        'graph', 'atomwise', 'atomwise', 'atomwise', 'atomwise', 'graph']\n",
    "ds_qm9.save_dataset('dataset_qm9', 1024, 128, 1024,\n",
    "                    norm_train=True, norm_valid=True, norm_test=True, mode=mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   ID.            Scale           Shift        Mode\n",
      "    0.     1.503183e+00    2.672827e+00       graph\n",
      "    1.     8.173763e+00    7.528103e+01       graph\n",
      "    2.     5.767056e+01   -6.306687e+02       graph\n",
      "    3.     1.229870e+02    3.108809e+01       graph\n",
      "    4.     1.238940e+02    6.617564e+02       graph\n",
      "    5.     2.804632e+02    1.189402e+03       graph\n",
      "    6.     8.699451e+01    3.914438e+02       graph\n",
      "    7.     1.824854e+01   -4.094204e+02    atomwise\n",
      "    8.     1.818079e+01   -4.118893e+02    atomwise\n",
      "    9.     1.816252e+01   -4.142276e+02    atomwise\n",
      "   10.     1.807273e+01   -3.811549e+02    atomwise\n",
      "   11.     6.082039e+00   -2.213512e+01       graph\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of Training, Validation and Test dataset: 1024, 128, 1024.\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Dataset file:\n",
      "   Training dataset        (origin):   dataset_qm9_origin_trainset_1024.npz\n",
      "   Validation dataset      (origin):   dataset_qm9_origin_validset_128.npz\n",
      "   Test dataset            (origin):   dataset_qm9_origin_testset_1024.npz\n",
      "=========================================================================================================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'name': 'qm9',\n",
       "  'num_atoms': 29,\n",
       "  'num_dataset': 1024,\n",
       "  'data_index': array([108284,  91708, 121372, ...,  16621,  36552,  50673], dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'type_ref': array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.31346686e+03, -1.30974915e+03,\n",
       "          -1.30726806e+03, -1.34143894e+03,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -9.93667075e+04, -9.93629871e+04,\n",
       "          -9.93605086e+04, -9.94048954e+04,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.43309938e+05, -1.43306220e+05,\n",
       "          -1.43303742e+05, -1.43349415e+05,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.97082067e+05, -1.97078349e+05,\n",
       "          -1.97075871e+05, -1.97121326e+05,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -2.61811546e+05, -2.61807828e+05,\n",
       "          -2.61805349e+05, -2.61850440e+05,  2.98100000e+00]]),\n",
       "  'Z': array([[6, 6, 6, ..., 0, 0, 0],\n",
       "         [8, 6, 6, ..., 0, 0, 0],\n",
       "         [8, 6, 6, ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [8, 6, 6, ..., 0, 0, 0]], dtype=int32),\n",
       "  'R': array([[[-1.1142875e-02,  1.5812907e-01,  4.4413228e-03],\n",
       "          [ 3.5342858e-03,  5.6812940e-03, -4.7988514e-03],\n",
       "          [ 7.2577603e-02, -5.4099612e-02,  1.1570221e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 6.1214771e-03,  1.4215238e-01, -1.3827609e-02],\n",
       "          [ 1.5041523e-03,  2.0087317e-03,  1.4307916e-03],\n",
       "          [ 6.6030197e-02, -5.1852573e-02,  1.3371411e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 2.0818852e-03,  1.4365010e-01, -2.3040319e-02],\n",
       "          [ 5.3538070e-03,  2.5866476e-03, -6.2697297e-03],\n",
       "          [-1.3854927e-01, -4.7261249e-02, -2.3530282e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[-1.0684761e-02,  1.5484959e-01, -1.2202791e-03],\n",
       "          [ 1.9219632e-03,  1.7879457e-03,  7.7857250e-05],\n",
       "          [ 7.9916753e-02, -5.3879716e-02, -1.1918602e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 1.5979379e-02, -1.2394467e-01, -3.5287064e-02],\n",
       "          [ 1.2903147e-02, -3.8187262e-03, -3.7626967e-02],\n",
       "          [ 1.2225902e-02,  1.4184269e-01, -3.7501145e-02],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[ 2.0368034e-02, -2.0850457e-02, -1.1839227e-02],\n",
       "          [ 1.0124510e-02,  9.8646425e-02, -1.7432189e-03],\n",
       "          [ 1.2764639e-01,  1.9380136e-01,  5.1811035e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]]], dtype=float32),\n",
       "  'E': array([[ 4.78800011e+00,  7.23499985e+01, -6.23293762e+02, ...,\n",
       "          -1.15278050e+06, -1.15289225e+06,  3.14960003e+01],\n",
       "         [ 3.38299990e+00,  7.03600006e+01, -6.37208923e+02, ...,\n",
       "          -1.15284000e+06, -1.15294375e+06,  2.86840000e+01],\n",
       "         [ 1.92770004e+00,  8.16600037e+01, -6.76591431e+02, ...,\n",
       "          -1.11694950e+06, -1.11707325e+06,  3.71080017e+01],\n",
       "         ...,\n",
       "         [ 1.35039997e+00,  7.68799973e+01, -5.83123596e+02, ...,\n",
       "          -8.60926500e+05, -8.61031812e+05,  3.00219994e+01],\n",
       "         [ 1.69309998e+00,  8.67500000e+01, -6.59525635e+02, ...,\n",
       "          -1.00974356e+06, -1.00987038e+06,  3.38540001e+01],\n",
       "         [ 4.14260006e+00,  7.58799973e+01, -6.80792175e+02, ...,\n",
       "          -1.05539425e+06, -1.05551100e+06,  3.11270008e+01]], dtype=float32),\n",
       "  'scale': array([  1.50318284,   8.17376272,  57.67055616, 122.98698488,\n",
       "         123.89396782, 280.4632002 ,  86.99450932,  18.24853761,\n",
       "          18.18078655,  18.16252178,  18.07272614,   6.08203921]),\n",
       "  'shift': array([   2.67282714,   75.28102688, -630.66868709,   31.08809144,\n",
       "          661.75640725, 1189.40203495,  391.44376141, -409.42037566,\n",
       "         -411.88933763, -414.22759985, -381.15492444,  -22.13511696]),\n",
       "  'scaleshift_mode': ('graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'graph'),\n",
       "  'use_atomwise_scaleshift': array([False, False, False, False, False, False, False,  True,  True,\n",
       "          True,  True, False])},\n",
       " {'name': 'qm9',\n",
       "  'num_atoms': 29,\n",
       "  'num_dataset': 128,\n",
       "  'data_index': array([119217,  83638,  17832,  73544,  72779, 133356,  52696,   4530,\n",
       "          48210,  66290,  88326, 126158,  64508,  22979,  85592,  52201,\n",
       "          38192,  13523,   5426,  41008, 121083,  37617,  58069, 127127,\n",
       "          92364,  65602, 133548,  50618,  91513, 107505,  12987,  39937,\n",
       "         123078,  44214,  39705,  95804,  70470, 100750,  33904, 116382,\n",
       "         109327,  51323,  71322,   2226,  68799,  44729,  52858,  94200,\n",
       "          51818,  10788, 102755,  62127,  36051,  22192,  26037,  50578,\n",
       "          37878, 117692, 128516,  14380,  56790, 108351, 116105,  75064,\n",
       "          43090, 115079,  77340,  15407, 113247, 105654,  88470, 112237,\n",
       "          39623,  34637, 132744,  16337, 130728,  76492, 115342,  83963,\n",
       "         113172, 109162, 111437,  12212, 122438, 130696, 126121,  12711,\n",
       "           7154,  64681,  50415,  64450,   8653, 132978,  14629,  51548,\n",
       "         105042,   2330,  20441,  62091,  40407,  97077,  47239,  41513,\n",
       "          95872, 103687,  80413,   2061,  32005, 108198,  97737, 113544,\n",
       "           3643,  81688, 117340,  11574,  47601, 125181, 119468, 110975,\n",
       "          53608,  32541,  59988,   6177,  44991, 109206, 128630,  21296],\n",
       "        dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'type_ref': array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.31346686e+03, -1.30974915e+03,\n",
       "          -1.30726806e+03, -1.34143894e+03,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -9.93667075e+04, -9.93629871e+04,\n",
       "          -9.93605086e+04, -9.94048954e+04,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.43309938e+05, -1.43306220e+05,\n",
       "          -1.43303742e+05, -1.43349415e+05,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.97082067e+05, -1.97078349e+05,\n",
       "          -1.97075871e+05, -1.97121326e+05,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -2.61811546e+05, -2.61807828e+05,\n",
       "          -2.61805349e+05, -2.61850440e+05,  2.98100000e+00]]),\n",
       "  'Z': array([[6, 8, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 8, ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [8, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 8, 6, ..., 0, 0, 0],\n",
       "         [6, 8, 6, ..., 0, 0, 0]], dtype=int32),\n",
       "  'R': array([[[-2.5081968e-02,  1.3415512e-01,  2.8292919e-04],\n",
       "          [-2.2156516e-02, -6.0766470e-03, -1.3487363e-02],\n",
       "          [ 7.7686816e-02, -5.0499283e-02, -1.0191814e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-6.0903680e-02,  1.3782467e-01,  9.7752549e-03],\n",
       "          [-5.2218386e-03, -3.0788200e-03, -4.9906294e-03],\n",
       "          [ 7.6597616e-02, -4.1184250e-02, -1.3205290e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-7.8710569e-03,  1.4489976e-01,  3.4626110e-03],\n",
       "          [-2.3963023e-03, -6.7769708e-03, -5.4659281e-04],\n",
       "          [ 9.2756651e-02, -7.0822537e-02,  3.8601402e-02],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[ 9.3458230e-03,  1.3364875e-01,  1.6986379e-02],\n",
       "          [ 1.0132672e-03, -4.4945139e-03, -4.9090753e-03],\n",
       "          [-1.4174549e-01, -5.1413324e-02, -8.1585599e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-3.5088016e-03,  1.3995786e-01,  4.4502923e-03],\n",
       "          [ 2.8237996e-03, -3.7284407e-03, -4.7976729e-03],\n",
       "          [ 6.2962823e-02, -5.3789377e-02, -1.1000546e-01],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]],\n",
       "  \n",
       "         [[-3.2752431e-03,  1.4069633e-01,  2.2779095e-03],\n",
       "          [-1.2417695e-03, -8.2001981e-04, -1.8107735e-02],\n",
       "          [-1.1851490e-01, -6.3027807e-02, -3.6884123e-03],\n",
       "          ...,\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01],\n",
       "          [-9.9900002e+01, -9.9900002e+01, -9.9900002e+01]]], dtype=float32),\n",
       "  'E': array([[ 3.8698001e+00,  7.8430000e+01, -7.0179620e+02, ...,\n",
       "          -1.1005364e+06, -1.1006586e+06,  3.5254002e+01],\n",
       "         [ 4.3329999e-01,  7.7419998e+01, -6.0911603e+02, ...,\n",
       "          -1.1136962e+06, -1.1138104e+06,  3.5266998e+01],\n",
       "         [ 3.3745999e+00,  7.0620003e+01, -6.1174152e+02, ...,\n",
       "          -9.5525688e+05, -9.5536619e+05,  2.8537001e+01],\n",
       "         ...,\n",
       "         [ 3.5392001e+00,  7.0160004e+01, -6.3589612e+02, ...,\n",
       "          -1.1950024e+06, -1.1951171e+06,  3.1047001e+01],\n",
       "         [ 4.3680000e+00,  6.0230000e+01, -7.4958032e+02, ...,\n",
       "          -1.3282602e+06, -1.3283691e+06,  2.6591000e+01],\n",
       "         [ 1.7701000e+00,  6.4830002e+01, -6.3852167e+02, ...,\n",
       "          -1.0496468e+06, -1.0497549e+06,  2.7480000e+01]], dtype=float32),\n",
       "  'scale': array([  1.50318284,   8.17376272,  57.67055616, 122.98698488,\n",
       "         123.89396782, 280.4632002 ,  86.99450932,  18.24853761,\n",
       "          18.18078655,  18.16252178,  18.07272614,   6.08203921]),\n",
       "  'shift': array([   2.67282714,   75.28102688, -630.66868709,   31.08809144,\n",
       "          661.75640725, 1189.40203495,  391.44376141, -409.42037566,\n",
       "         -411.88933763, -414.22759985, -381.15492444,  -22.13511696]),\n",
       "  'scaleshift_mode': ('graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'graph'),\n",
       "  'use_atomwise_scaleshift': array([False, False, False, False, False, False, False,  True,  True,\n",
       "          True,  True, False])},\n",
       " {'name': 'qm9',\n",
       "  'num_atoms': 29,\n",
       "  'num_dataset': 1024,\n",
       "  'data_index': array([  9948, 110902,  41365, ...,  18549,  78512,  74618], dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'type_ref': array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.31346686e+03, -1.30974915e+03,\n",
       "          -1.30726806e+03, -1.34143894e+03,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -9.93667075e+04, -9.93629871e+04,\n",
       "          -9.93605086e+04, -9.94048954e+04,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.43309938e+05, -1.43306220e+05,\n",
       "          -1.43303742e+05, -1.43349415e+05,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -1.97082067e+05, -1.97078349e+05,\n",
       "          -1.97075871e+05, -1.97121326e+05,  2.98100000e+00],\n",
       "         [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
       "           0.00000000e+00, -2.61811546e+05, -2.61807828e+05,\n",
       "          -2.61805349e+05, -2.61850440e+05,  2.98100000e+00]]),\n",
       "  'Z': array([[6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [8, 6, 7, ..., 0, 0, 0],\n",
       "         [6, 6, 8, ..., 0, 0, 0],\n",
       "         [6, 6, 6, ..., 0, 0, 0]], dtype=int32),\n",
       "  'R': array([[[-1.62902959e-02,  1.36730880e-01, -9.98675730e-03],\n",
       "          [ 1.38593940e-02, -1.06998775e-02,  4.15915297e-03],\n",
       "          [ 1.52657643e-01, -6.42142296e-02, -8.41158908e-03],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[ 1.10524502e-02,  1.58110440e-01,  1.08578904e-02],\n",
       "          [-4.21275035e-04,  5.68928989e-03,  4.74962144e-04],\n",
       "          [ 6.19037785e-02, -5.18585853e-02, -1.26719937e-01],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[-8.54106329e-04,  1.45507395e-01, -3.53108719e-02],\n",
       "          [ 1.44421622e-01,  1.60085946e-01,  2.47270558e-02],\n",
       "          [ 1.42592028e-01,  8.34128074e-03,  5.83807863e-02],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[-9.12077143e-04, -1.42512412e-03, -8.75520986e-03],\n",
       "          [ 5.50002500e-04,  1.18384898e-01, -2.63127428e-03],\n",
       "          [-9.84417796e-02,  2.13059574e-01,  8.81935563e-03],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[ 6.42315764e-03,  1.41510636e-01, -1.73014812e-02],\n",
       "          [ 3.31989228e-04, -6.93340274e-03, -2.68555875e-03],\n",
       "          [ 5.51430434e-02, -9.14408639e-02, -1.08960524e-01],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]],\n",
       "  \n",
       "         [[-2.74167373e-03,  1.50980234e-01,  2.23346078e-03],\n",
       "          [-1.16161034e-02,  1.88250549e-03, -3.42915463e-03],\n",
       "          [-1.13114680e-03, -7.84160644e-02, -1.10561773e-01],\n",
       "          ...,\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01],\n",
       "          [-9.99000015e+01, -9.99000015e+01, -9.99000015e+01]]],\n",
       "        dtype=float32),\n",
       "  'E': array([[ 2.64450002e+00,  6.93499985e+01, -6.49548767e+02, ...,\n",
       "          -1.01063650e+06, -1.01075212e+06,  3.33450012e+01],\n",
       "         [ 1.89260006e+00,  7.73199997e+01, -6.39834412e+02, ...,\n",
       "          -1.11373250e+06, -1.11384238e+06,  3.29049988e+01],\n",
       "         [ 1.53890002e+00,  7.34800034e+01, -6.33008118e+02, ...,\n",
       "          -1.05523875e+06, -1.05534062e+06,  2.80559998e+01],\n",
       "         ...,\n",
       "         [ 3.72429991e+00,  6.63300018e+01, -6.14629578e+02, ...,\n",
       "          -9.52241875e+05, -9.52341438e+05,  2.47129993e+01],\n",
       "         [ 2.29900002e+00,  7.57099991e+01, -5.84961426e+02, ...,\n",
       "          -1.11042412e+06, -1.11053462e+06,  3.37949982e+01],\n",
       "         [ 1.32099998e+00,  7.99499969e+01, -5.67370605e+02, ...,\n",
       "          -1.01305244e+06, -1.01315881e+06,  3.17670002e+01]], dtype=float32),\n",
       "  'scale': array([  1.50318284,   8.17376272,  57.67055616, 122.98698488,\n",
       "         123.89396782, 280.4632002 ,  86.99450932,  18.24853761,\n",
       "          18.18078655,  18.16252178,  18.07272614,   6.08203921]),\n",
       "  'shift': array([   2.67282714,   75.28102688, -630.66868709,   31.08809144,\n",
       "          661.75640725, 1189.40203495,  391.44376141, -409.42037566,\n",
       "         -411.88933763, -414.22759985, -381.15492444,  -22.13511696]),\n",
       "  'scaleshift_mode': ('graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'graph',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'atomwise',\n",
       "   'graph'),\n",
       "  'use_atomwise_scaleshift': array([False, False, False, False, False, False, False,  True,  True,\n",
       "          True,  True, False])})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_qm9.save_dataset('dataset_qm9', 1024, 128, 1024, norm_train=False,\n",
    "                    norm_valid=False, norm_test=False, mode=mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E (555092, 1)\n",
      "name ()\n",
      "F (555092, 9, 3)\n",
      "theory ()\n",
      "R (555092, 9, 3)\n",
      "z (9,)\n",
      "type ()\n",
      "md5 ()\n"
     ]
    }
   ],
   "source": [
    "md17_data = np.load('ethanol_dft.npz')\n",
    "for k, v in md17_data.items():\n",
    "    print(k, v.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========================================================================================================\n",
      "Number of atoms: 9\n",
      "Number of data: 555092\n",
      "Dimension of space: 3\n",
      "Total number of effective atoms: 4995828\n",
      "Shape of atom types (Z): (1, 9)\n",
      "Shape of position (R): (555092, 9, 3)\n",
      "Shape of label (E): (555092, 1)\n",
      "Shape of force (F): (555092, 9, 3)\n",
      "Length unit: Angstrom\n",
      "Energy unit: kcal mol-1\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to graph:\n",
      "               Mean             STD             Min             Max          Median           Range\n",
      "      -9.719593e+04    4.192011e+00   -9.720841e+04   -9.717179e+04   -9.719637e+04    3.661111e+01\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to each atom:\n",
      "               Mean             STD             Min             Max          Median           Range\n",
      "      -1.079955e+04    4.657790e-01   -1.080093e+04   -1.079687e+04   -1.079960e+04    4.067901e+00\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Force norm analysis:\n",
      "               Mean             STD             Max          Median\n",
      "       3.899777e+01    2.339806e+01    2.435161e+02    3.381680e+01\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of data: 555092\n",
      "=========================================================================================================\n"
     ]
    }
   ],
   "source": [
    "ds_md17 = DatasetProcessor(\n",
    "    name='md17',\n",
    "    atom_types=md17_data['z'],\n",
    "    position=md17_data['R'],\n",
    "    label=md17_data['E'],\n",
    "    force=md17_data['F'],\n",
    "    length_unit='A',\n",
    "    energy_unit='kcal/mol',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Change the length unit from \"Angstrom\" to \"nm\"\n",
      "Change the energy unit from \"kcal mol-1\" to \"kJ mol-1\"\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to graph:\n",
      "               Mean             STD             Min             Max          Median           Range\n",
      "      -4.066678e+05    1.753937e+01   -4.067200e+05   -4.065668e+05   -4.066696e+05    1.531809e+02\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Label analysis according to each atom:\n",
      "               Mean             STD             Min             Max          Median           Range\n",
      "      -4.518531e+04    1.948819e+00   -4.519111e+04   -4.517409e+04   -4.518551e+04    1.702010e+01\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Force norm analysis:\n",
      "               Mean             STD             Max          Median\n",
      "       1.631667e+03    9.789750e+02    1.018871e+04    1.414895e+03\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of data: 555092\n",
      "=========================================================================================================\n"
     ]
    }
   ],
   "source": [
    "ds_md17.convert_units('nm', 'kj/mol')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([378956, 348096, 112136, ..., 185742,  23029, 205694], dtype=int32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_md17.shuffle_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              Scale           Shift        Mode\n",
      "       1.948819e+00   -4.518531e+04    atomwise\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Number of Training, Validation and Test dataset: 1024, 128, 1024.\n",
      "---------------------------------------------------------------------------------------------------------\n",
      "Dataset file:\n",
      "   Training dataset    (normalized):   dataset_ethanol_normed_trainset_1024.npz\n",
      "   Validation dataset  (normalized):   dataset_ethanol_normed_validset_128.npz\n",
      "   Test dataset            (origin):   dataset_ethanol_origin_testset_1024.npz\n",
      "=========================================================================================================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'name': 'md17',\n",
       "  'num_atoms': 9,\n",
       "  'num_dataset': 1024,\n",
       "  'data_index': array([378956, 348096, 112136, ...,  36787, 295704, 409501], dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'Z': array([[6, 6, 8, 1, 1, 1, 1, 1, 1]], dtype=int32),\n",
       "  'R': array([[[ 0.03992085,  0.03301192,  0.02689481],\n",
       "          [-0.0639054 ,  0.09025992, -0.06854789],\n",
       "          [ 0.0276721 , -0.11235844,  0.0422143 ],\n",
       "          ...,\n",
       "          [-0.16907883,  0.08484092, -0.03171887],\n",
       "          [-0.03656308,  0.19641401, -0.09884108],\n",
       "          [-0.05365539, -0.13611984, -0.00298879]],\n",
       "  \n",
       "         [[ 0.03246842, -0.0045715 , -0.04411067],\n",
       "          [-0.12228684,  0.014885  , -0.04969889],\n",
       "          [ 0.07969039, -0.01376768,  0.08598094],\n",
       "          ...,\n",
       "          [-0.16438967, -0.02044584, -0.14762333],\n",
       "          [-0.1743204 , -0.03800723,  0.04264783],\n",
       "          [ 0.14583254,  0.05424568,  0.11390269]],\n",
       "  \n",
       "         [[-0.03228785, -0.01766151, -0.03811911],\n",
       "          [-0.08749443,  0.05696383,  0.07786967],\n",
       "          [ 0.10875645, -0.03407211, -0.04460916],\n",
       "          ...,\n",
       "          [-0.05604751,  0.00498362,  0.17511088],\n",
       "          [-0.05017642,  0.15849467,  0.11213855],\n",
       "          [ 0.13986874, -0.04554918,  0.04901359]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[ 0.00192143,  0.00532413,  0.05452912],\n",
       "          [-0.11617309,  0.05190711, -0.02985357],\n",
       "          [ 0.1023763 , -0.04852812, -0.02846017],\n",
       "          ...,\n",
       "          [-0.09372877,  0.13768572, -0.08306124],\n",
       "          [-0.15245484, -0.03122399, -0.09280465],\n",
       "          [ 0.16813427, -0.08905247,  0.03366383]],\n",
       "  \n",
       "         [[ 0.02702162, -0.04376558, -0.03860364],\n",
       "          [-0.13103542,  0.00438406, -0.03872851],\n",
       "          [ 0.09743995,  0.02889462,  0.07182059],\n",
       "          ...,\n",
       "          [-0.19674376, -0.05705723, -0.10028627],\n",
       "          [-0.17743057,  0.01592045,  0.05908336],\n",
       "          [ 0.07673421,  0.11911692,  0.05954019]],\n",
       "  \n",
       "         [[ 0.03761427,  0.03562684,  0.02659132],\n",
       "          [-0.06197226,  0.08860236, -0.07310329],\n",
       "          [ 0.02057528, -0.11160219,  0.0381396 ],\n",
       "          ...,\n",
       "          [-0.15019402,  0.02041352, -0.08156363],\n",
       "          [-0.09681427,  0.18423726, -0.03742639],\n",
       "          [ 0.06911577, -0.14681178,  0.10825907]]], dtype=float32),\n",
       "  'E': array([[-8.962144 ],\n",
       "         [-4.1365805],\n",
       "         [-2.1421676],\n",
       "         ...,\n",
       "         [-9.710595 ],\n",
       "         [-6.200368 ],\n",
       "         [-4.3626637]], dtype=float32),\n",
       "  'F': array([[[  737.8506   ,   -33.839344 ,  -361.60486  ],\n",
       "          [  -76.06496  ,   612.82     ,   285.4191   ],\n",
       "          [  411.30368  ,   491.872    ,   -19.966518 ],\n",
       "          ...,\n",
       "          [  256.19925  ,   183.27121  ,   176.5632   ],\n",
       "          [ -235.77473  ,  -500.90445  ,   168.12238  ],\n",
       "          [ -233.34729  ,  -134.86061  ,  -151.46947  ]],\n",
       "  \n",
       "         [[ -665.40375  ,  -537.8482   ,  -291.89886  ],\n",
       "          [  -52.75404  , -1123.9679   ,   270.85895  ],\n",
       "          [  396.76035  ,    39.274715 ,  1002.67535  ],\n",
       "          ...,\n",
       "          [  135.79831  ,  -201.72478  ,   350.7262   ],\n",
       "          [  436.88956  ,   505.62762  ,  -922.4308   ],\n",
       "          [ -350.8861   ,  -222.74341  ,  -369.6001   ]],\n",
       "  \n",
       "         [[  -48.442036 , -1087.2024   ,  -227.14577  ],\n",
       "          [  561.78595  ,   507.3258   ,  1397.4519   ],\n",
       "          [  328.906    ,    63.33812  ,   692.10406  ],\n",
       "          ...,\n",
       "          [ -334.74115  ,   -74.723755 ,  -494.75443  ],\n",
       "          [   43.678112 ,  -207.94708  ,  -707.3281   ],\n",
       "          [  -55.17504  ,    -5.2979803,  -419.59152  ]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[ -519.1932   ,  -768.1502   ,    67.13498  ],\n",
       "          [ -451.39972  , -1120.6656   ,  1600.603    ],\n",
       "          [  562.6375   ,  -478.43372  ,   137.52568  ],\n",
       "          ...,\n",
       "          [   53.428356 ,  1012.9089   ,  -614.42413  ],\n",
       "          [   82.69473  ,   128.32935  ,    -5.324715 ],\n",
       "          [ -171.87567  ,   146.75816  ,  -385.90744  ]],\n",
       "  \n",
       "         [[ -646.10504  ,   177.99353  ,  1086.94     ],\n",
       "          [ 1123.7649   ,   450.97592  ,  -857.7322   ],\n",
       "          [  -50.59247  , -1427.6475   ,  -368.07126  ],\n",
       "          ...,\n",
       "          [   23.296211 ,  -231.21542  ,   -88.17738  ],\n",
       "          [  164.12897  ,   -15.320773 ,   181.59631  ],\n",
       "          [ -213.75996  ,  1088.8015   ,  -135.00023  ]],\n",
       "  \n",
       "         [[  838.19586  ,  -468.46652  ,   887.8585   ],\n",
       "          [ -850.90485  ,  -473.51596  ,   966.87524  ],\n",
       "          [ -788.2559   ,   991.2189   , -1384.1686   ],\n",
       "          ...,\n",
       "          [  126.6401   ,   147.55724  ,   336.97205  ],\n",
       "          [ -159.55283  ,   329.99173  ,   -45.78443  ],\n",
       "          [  831.022    ,  -359.81757  ,  1094.0891   ]]], dtype=float32),\n",
       "  'scale': array([1.9488194], dtype=float32),\n",
       "  'shift': array([-45185.31], dtype=float32),\n",
       "  'scaleshift_mode': 'atomwise',\n",
       "  'use_atomwise_scaleshift': array([ True]),\n",
       "  'avg_force_dis': array([0.00119437], dtype=float32)},\n",
       " {'name': 'md17',\n",
       "  'num_atoms': 9,\n",
       "  'num_dataset': 128,\n",
       "  'data_index': array([ 87137,  52497,  41760, 133931, 351197,  37597, 144680, 246842,\n",
       "          18063, 489922,  84887, 281286, 550016, 277897, 220236,  86991,\n",
       "         507881,   6894, 214206, 372142,  30377, 164872,  62087, 336444,\n",
       "         537204, 471325,  88469, 544449,  35771, 204340, 194089, 512117,\n",
       "         287937, 210100, 326052,  24372, 257700,  27364, 129466, 403057,\n",
       "          74279, 104084, 519162, 180792, 498291, 452882, 149663,  23388,\n",
       "         347566, 102345, 350320, 106591, 535352, 530088, 260663, 157855,\n",
       "         525734, 302549, 184408, 373211, 485147, 540804, 325609, 529855,\n",
       "         297472, 418295,  84471, 298390,  31994, 251533, 519589, 272292,\n",
       "         487812, 403586,  70641, 405161, 510628, 536976, 478643, 374266,\n",
       "         171393,   6783,  80630, 308810, 123789, 390430, 387378, 436828,\n",
       "         266384, 413919, 430450, 498190, 470228, 232219, 412374, 266177,\n",
       "         501379, 338790, 504918, 164739,  49561, 385393, 155582, 178995,\n",
       "         157596, 250423, 157625,  15240, 290611,   2646, 288482,  19675,\n",
       "         162518, 281640, 309951, 123952, 531573, 538487, 321767, 334552,\n",
       "          68535,  76632, 351588, 465014, 465808, 354287, 348260,  91762],\n",
       "        dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'Z': array([[6, 6, 8, 1, 1, 1, 1, 1, 1]], dtype=int32),\n",
       "  'R': array([[[-7.33467005e-03, -4.94685359e-02,  3.09051964e-02],\n",
       "          [-1.00180790e-01,  7.37974495e-02,  3.82807739e-02],\n",
       "          [ 9.63298306e-02, -2.15023104e-02, -6.50609508e-02],\n",
       "          ...,\n",
       "          [-1.57053068e-01,  7.09956437e-02,  1.27748892e-01],\n",
       "          [-5.33375107e-02,  1.71536669e-01,  2.87328716e-02],\n",
       "          [ 1.80010468e-01, -5.06403707e-02, -2.87131425e-02]],\n",
       "  \n",
       "         [[-2.96711959e-02, -4.68306094e-02, -8.21130991e-04],\n",
       "          [-8.99534449e-02,  9.26802382e-02,  2.30049435e-02],\n",
       "          [ 1.07694529e-01, -3.85197103e-02, -2.14604475e-02],\n",
       "          ...,\n",
       "          [-2.70249415e-02,  1.33143261e-01,  1.04822636e-01],\n",
       "          [-9.02404636e-02,  1.63581073e-01, -6.23214841e-02],\n",
       "          [ 1.43024608e-01, -1.18374258e-01, -4.94678319e-02]],\n",
       "  \n",
       "         [[-9.31464788e-03, -3.09217479e-02,  4.80818190e-02],\n",
       "          [-1.04140013e-01,  7.76589811e-02, -1.89293269e-02],\n",
       "          [ 1.02793843e-01, -5.07492013e-02, -2.78528687e-02],\n",
       "          ...,\n",
       "          [-4.93561141e-02,  1.72942281e-01, -3.93252932e-02],\n",
       "          [-1.27391040e-01,  5.15612885e-02, -1.16997950e-01],\n",
       "          [ 1.40150949e-01,  3.36055867e-02,  1.14980003e-05]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[ 4.35832292e-02,  3.84615622e-02,  5.29069314e-03],\n",
       "          [-8.12560841e-02,  6.10200129e-02, -8.43458474e-02],\n",
       "          [ 4.24023047e-02, -8.79103690e-02,  7.30142295e-02],\n",
       "          ...,\n",
       "          [-1.82480365e-01,  1.38873281e-02, -3.83037850e-02],\n",
       "          [-1.02627948e-01,  1.65212959e-01, -1.01515248e-01],\n",
       "          [-3.20426896e-02, -1.47170201e-01,  6.58851489e-02]],\n",
       "  \n",
       "         [[ 2.53772158e-02, -2.81953402e-02, -4.71129008e-02],\n",
       "          [-1.23050474e-01,  2.11861022e-02, -5.43271154e-02],\n",
       "          [ 9.38291401e-02,  4.43330221e-03,  8.59732479e-02],\n",
       "          ...,\n",
       "          [-1.31901056e-01,  7.25662857e-02, -1.53461650e-01],\n",
       "          [-1.97769761e-01, -5.17504252e-02, -4.66614291e-02],\n",
       "          [ 3.85163128e-02, -1.30492542e-02,  1.70459867e-01]],\n",
       "  \n",
       "         [[-2.55391672e-02, -4.64483686e-02,  7.10871723e-03],\n",
       "          [-9.38182622e-02,  9.05697122e-02,  3.50831933e-02],\n",
       "          [ 1.12512723e-01, -4.18050885e-02, -3.18864435e-02],\n",
       "          ...,\n",
       "          [-1.41935218e-02,  1.66850016e-01,  3.38415094e-02],\n",
       "          [-1.60126388e-01,  1.08510673e-01, -6.31183982e-02],\n",
       "          [ 9.91470218e-02, -2.50386354e-02, -1.26092196e-01]]],\n",
       "        dtype=float32),\n",
       "  'E': array([[ 6.99750376e+00],\n",
       "         [-8.26256466e+00],\n",
       "         [ 7.57142639e+00],\n",
       "         [-5.58261299e+00],\n",
       "         [ 1.67216797e+01],\n",
       "         [ 4.54604626e-01],\n",
       "         [ 2.16338563e+00],\n",
       "         [-6.83441639e+00],\n",
       "         [ 6.69511223e+00],\n",
       "         [-2.88113070e+00],\n",
       "         [-5.47643471e+00],\n",
       "         [ 1.50258434e+00],\n",
       "         [-4.97241735e+00],\n",
       "         [ 1.05174589e+01],\n",
       "         [-5.58147526e+00],\n",
       "         [-7.29539824e+00],\n",
       "         [ 4.93723917e+00],\n",
       "         [-7.90347242e+00],\n",
       "         [ 2.06635342e+01],\n",
       "         [-3.78215909e+00],\n",
       "         [ 6.42431676e-01],\n",
       "         [ 5.13126373e+00],\n",
       "         [-1.24592447e+01],\n",
       "         [ 1.06418724e+01],\n",
       "         [-9.34060001e+00],\n",
       "         [-1.09566379e+00],\n",
       "         [-9.14070034e+00],\n",
       "         [-2.44306588e+00],\n",
       "         [-2.58742285e+00],\n",
       "         [ 5.33272552e+00],\n",
       "         [-2.62201548e+00],\n",
       "         [-1.71528664e+01],\n",
       "         [ 7.15062904e+00],\n",
       "         [ 2.10844383e+01],\n",
       "         [ 4.96852779e+00],\n",
       "         [ 6.23075962e+00],\n",
       "         [ 8.43557250e-03],\n",
       "         [-9.54394341e+00],\n",
       "         [-7.36456966e+00],\n",
       "         [-9.79152679e+00],\n",
       "         [-7.85445261e+00],\n",
       "         [ 1.60479126e+01],\n",
       "         [-4.19707537e+00],\n",
       "         [-6.30103207e+00],\n",
       "         [ 2.62956929e+00],\n",
       "         [-5.81313944e+00],\n",
       "         [ 4.31632757e+00],\n",
       "         [-1.27343235e+01],\n",
       "         [-4.29099464e+00],\n",
       "         [ 2.25139880e+00],\n",
       "         [-1.12866964e+01],\n",
       "         [ 1.16671658e+01],\n",
       "         [ 1.36348426e+00],\n",
       "         [-6.01437426e+00],\n",
       "         [-2.19401026e+00],\n",
       "         [ 6.41177034e+00],\n",
       "         [-8.45654726e-01],\n",
       "         [-1.44840612e+01],\n",
       "         [ 1.93365827e-01],\n",
       "         [ 8.64059353e+00],\n",
       "         [-5.87803030e+00],\n",
       "         [-3.09639192e+00],\n",
       "         [ 9.83612001e-01],\n",
       "         [ 9.80125237e+00],\n",
       "         [ 6.30527830e+00],\n",
       "         [ 6.92723417e+00],\n",
       "         [-9.55769539e+00],\n",
       "         [-9.67041206e+00],\n",
       "         [ 4.93403769e+00],\n",
       "         [-1.14396453e+00],\n",
       "         [-6.97096884e-01],\n",
       "         [ 1.08060102e+01],\n",
       "         [ 4.29708862e+00],\n",
       "         [-1.70524883e+01],\n",
       "         [-2.38101053e+00],\n",
       "         [ 1.43533287e+01],\n",
       "         [-5.10931826e+00],\n",
       "         [-1.40365076e+00],\n",
       "         [-2.59879208e+00],\n",
       "         [-7.35288143e+00],\n",
       "         [ 3.72232819e+00],\n",
       "         [ 1.52393122e+01],\n",
       "         [-1.34938478e+01],\n",
       "         [ 9.45750237e+00],\n",
       "         [-3.19639158e+00],\n",
       "         [ 7.05275393e+00],\n",
       "         [-1.38751955e+01],\n",
       "         [ 5.21261692e+00],\n",
       "         [ 1.22059381e+00],\n",
       "         [-1.46102047e+00],\n",
       "         [ 1.56401587e+00],\n",
       "         [-6.17933178e+00],\n",
       "         [ 2.23875313e+01],\n",
       "         [ 1.74504387e+00],\n",
       "         [ 9.75866318e+00],\n",
       "         [-2.26659226e+00],\n",
       "         [-6.90827990e+00],\n",
       "         [-1.67218246e+01],\n",
       "         [-1.03244555e+00],\n",
       "         [-2.31952047e+00],\n",
       "         [ 6.83996820e+00],\n",
       "         [ 1.11541576e+01],\n",
       "         [-1.42760820e+01],\n",
       "         [ 3.28317356e+00],\n",
       "         [-1.27046833e+01],\n",
       "         [-9.22272682e-01],\n",
       "         [-9.85194778e+00],\n",
       "         [ 1.92467060e+01],\n",
       "         [-1.15775478e+00],\n",
       "         [ 3.05412960e+00],\n",
       "         [-3.65199161e+00],\n",
       "         [-1.01590929e+01],\n",
       "         [ 1.62069893e+01],\n",
       "         [ 2.86698389e+00],\n",
       "         [ 1.47775993e+01],\n",
       "         [-2.86632752e+00],\n",
       "         [ 1.97969615e+00],\n",
       "         [ 1.12557030e+01],\n",
       "         [ 3.28001785e+00],\n",
       "         [ 5.94807267e-01],\n",
       "         [-5.24057913e+00],\n",
       "         [ 1.06622303e+00],\n",
       "         [-1.64343524e+00],\n",
       "         [ 5.37658453e+00],\n",
       "         [-7.58928585e+00],\n",
       "         [ 7.03452158e+00],\n",
       "         [ 1.86373253e+01],\n",
       "         [ 1.30922794e+01]], dtype=float32),\n",
       "  'F': array([[[  154.92223  ,  -277.27286  ,    81.25798  ],\n",
       "          [-1177.8884   ,  -230.61983  , -1101.7235   ],\n",
       "          [ -199.87857  ,   102.797455 ,  -185.45308  ],\n",
       "          ...,\n",
       "          [ -251.33994  ,    23.998005 ,   720.7105   ],\n",
       "          [  249.07884  ,    14.886973 ,    19.24829  ],\n",
       "          [  255.50117  ,  -120.53116  ,    90.01071  ]],\n",
       "  \n",
       "         [[ -541.4228   ,   154.43655  ,   627.6629   ],\n",
       "          [ 1123.2649   ,    47.53993  ,  -107.563705 ],\n",
       "          [  146.12395  ,  1668.5198   ,  -194.48752  ],\n",
       "          ...,\n",
       "          [ -148.44742  ,   185.69765  ,  -129.60768  ],\n",
       "          [   66.371056 ,  -210.15767  ,    66.083405 ],\n",
       "          [  529.7584   , -1404.6356   ,  -471.21628  ]],\n",
       "  \n",
       "         [[-1520.0416   ,   955.1118   ,    44.044876 ],\n",
       "          [  264.6069   ,   167.79018  ,  1260.9896   ],\n",
       "          [  796.29126  ,    78.55527  ,   -98.49411  ],\n",
       "          ...,\n",
       "          [  -54.699753 ,   -87.145615 ,   405.9854   ],\n",
       "          [ -610.65155  ,  -513.58374  ,  -905.86505  ],\n",
       "          [  852.6228   ,   160.0499   ,  -462.24768  ]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[-2125.663    ,   313.62946  ,  1310.2463   ],\n",
       "          [ -402.91187  , -1210.0991   ,  -441.60876  ],\n",
       "          [ 1047.902    ,   -87.859535 ,    87.90657  ],\n",
       "          ...,\n",
       "          [ 1110.5311   ,   428.7235   ,  -433.86502  ],\n",
       "          [    8.701442 ,   372.90384  ,   -10.961289 ],\n",
       "          [ -481.9751   ,    62.551914 ,  -351.6624   ]],\n",
       "  \n",
       "         [[ -401.19885  ,   662.01733  ,  1266.5792   ],\n",
       "          [  782.5227   ,  1356.2601   ,   762.9617   ],\n",
       "          [ -684.76184  ,  -443.89658  ,   351.72726  ],\n",
       "          ...,\n",
       "          [ -182.19092  ,    27.711046 ,   405.2591   ],\n",
       "          [ -545.6509   ,  -790.0778   ,    30.460886 ],\n",
       "          [  341.0318   ,    -2.8380113, -1013.6812   ]],\n",
       "  \n",
       "         [[-1283.239    ,    32.352932 ,  -749.09906  ],\n",
       "          [ -145.21793  ,    25.146866 ,  -934.0011   ],\n",
       "          [ -669.47833  ,   -87.374344 ,  -154.15088  ],\n",
       "          ...,\n",
       "          [ -105.23083  ,    61.87364  ,    99.62094  ],\n",
       "          [  403.87927  ,   110.53525  ,  1235.7499   ],\n",
       "          [  423.3374   ,    89.4676   ,  -282.91443  ]]], dtype=float32),\n",
       "  'scale': array([1.9488194], dtype=float32),\n",
       "  'shift': array([-45185.31], dtype=float32),\n",
       "  'scaleshift_mode': 'atomwise',\n",
       "  'use_atomwise_scaleshift': array([ True]),\n",
       "  'avg_force_dis': array([0.00119437], dtype=float32)},\n",
       " {'name': 'md17',\n",
       "  'num_atoms': 9,\n",
       "  'num_dataset': 1024,\n",
       "  'data_index': array([460728, 421500, 233686, ..., 470372, 387294, 510807], dtype=int32),\n",
       "  'normed_dataset': False,\n",
       "  'length_unit': 'nm',\n",
       "  'energy_unit': 'kj/mol',\n",
       "  'Z': array([[6, 6, 8, 1, 1, 1, 1, 1, 1]], dtype=int32),\n",
       "  'R': array([[[ 3.15719247e-02,  3.76577824e-02,  1.79349240e-02],\n",
       "          [-6.47983477e-02,  8.61792192e-02, -8.60510468e-02],\n",
       "          [ 2.38753241e-02, -1.13392726e-01,  6.57409132e-02],\n",
       "          ...,\n",
       "          [-7.88076967e-02,  9.39244870e-03, -1.71162024e-01],\n",
       "          [-1.59875080e-01,  1.23354256e-01, -5.24685346e-02],\n",
       "          [ 1.18793011e-01, -1.35041341e-01,  6.80906102e-02]],\n",
       "  \n",
       "         [[ 3.08904201e-02,  3.22379954e-02,  4.56425212e-02],\n",
       "          [-5.16040064e-02,  8.83255526e-02, -7.99379945e-02],\n",
       "          [ 1.84098296e-02, -1.09041616e-01,  4.09286357e-02],\n",
       "          ...,\n",
       "          [-1.55658379e-01,  8.06336924e-02, -8.01753551e-02],\n",
       "          [-1.74019746e-02,  1.96167126e-01, -1.05445690e-01],\n",
       "          [ 4.23220769e-02, -1.30937934e-01, -5.21939658e-02]],\n",
       "  \n",
       "         [[-2.48317397e-03, -1.26837566e-02, -4.89062741e-02],\n",
       "          [-1.27857536e-01, -3.78823765e-02,  3.72334272e-02],\n",
       "          [ 1.17048740e-01,  5.19638509e-02,  8.81022215e-03],\n",
       "          ...,\n",
       "          [-1.08718082e-01, -1.29996851e-01,  9.57021490e-02],\n",
       "          [-1.33644372e-01,  5.57788908e-02,  9.19445530e-02],\n",
       "          [ 1.67279258e-01, -1.80356540e-02,  4.80102040e-02]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[ 3.34200263e-02, -2.16021389e-03, -3.71256731e-02],\n",
       "          [-5.32608703e-02,  1.11915775e-01, -2.62740608e-02],\n",
       "          [ 2.54980251e-02, -1.00974910e-01,  5.65546677e-02],\n",
       "          ...,\n",
       "          [-1.52919367e-01,  7.35318959e-02, -3.45439836e-02],\n",
       "          [-4.73069847e-02,  1.57425836e-01,  7.03847259e-02],\n",
       "          [-5.68791889e-02, -7.51329660e-02,  1.01453073e-01]],\n",
       "  \n",
       "         [[-1.68035477e-02, -8.46538134e-03, -5.28748557e-02],\n",
       "          [-6.48101093e-03,  1.31398529e-01, -1.58145995e-04],\n",
       "          [ 1.65158678e-02, -1.16859555e-01,  4.58596349e-02],\n",
       "          ...,\n",
       "          [ 9.55277309e-02,  1.68825030e-01,  1.94089860e-02],\n",
       "          [-3.02045625e-02,  2.09748209e-01, -7.47435838e-02],\n",
       "          [ 1.04158230e-01, -9.14058164e-02,  7.72011653e-02]],\n",
       "  \n",
       "         [[-2.98693124e-02,  8.92487995e-04, -4.52983528e-02],\n",
       "          [ 9.23815668e-02,  8.22490752e-02, -3.56875286e-02],\n",
       "          [-6.02385141e-02, -7.10020736e-02,  7.83020630e-02],\n",
       "          ...,\n",
       "          [ 1.00357249e-01,  1.18727475e-01, -1.38685450e-01],\n",
       "          [ 1.12768471e-01,  1.63158417e-01,  3.52819227e-02],\n",
       "          [-1.89843457e-02, -1.51558235e-01,  6.32585362e-02]]],\n",
       "        dtype=float32),\n",
       "  'E': array([[-406639.25],\n",
       "         [-406669.7 ],\n",
       "         [-406675.6 ],\n",
       "         ...,\n",
       "         [-406652.9 ],\n",
       "         [-406679.66],\n",
       "         [-406660.5 ]], dtype=float32),\n",
       "  'F': array([[[ -155.59187 , -1913.2684  ,   -66.58135 ],\n",
       "          [  -52.281673,  -252.09567 , -2852.738   ],\n",
       "          [   74.95872 ,  2530.0264  , -1301.9137  ],\n",
       "          ...,\n",
       "          [ -134.25722 ,   479.7451  ,  1384.1571  ],\n",
       "          [ -773.2394  ,   -95.80216 ,   309.5059  ],\n",
       "          [  -56.053543,  -218.74194 ,    42.77843 ]],\n",
       "  \n",
       "         [[-3334.8616  ,  -148.31598 , -1929.6322  ],\n",
       "          [ 3626.3025  ,   622.40704 ,  -895.59424 ],\n",
       "          [  936.7864  ,  -701.03906 ,  -705.0769  ],\n",
       "          ...,\n",
       "          [-1934.6205  ,   -87.2629  ,   506.15482 ],\n",
       "          [ -636.63275 , -1345.3606  ,   344.66968 ],\n",
       "          [ -120.65387 ,  -329.94034 ,   547.96155 ]],\n",
       "  \n",
       "         [[  -83.90408 ,  -527.7992  , -2337.945   ],\n",
       "          [ 1117.966   ,  -508.126   ,   522.3276  ],\n",
       "          [-2253.222   ,   607.6256  ,  -589.2203  ],\n",
       "          ...,\n",
       "          [  120.48792 ,   624.0039  ,   714.95306 ],\n",
       "          [ -684.62317 ,   -56.701725,   857.56586 ],\n",
       "          [  719.01794 ,  -942.808   ,   648.33826 ]],\n",
       "  \n",
       "         ...,\n",
       "  \n",
       "         [[  424.00598 , -2263.98    , -4059.624   ],\n",
       "          [-1241.553   ,   119.3687  ,  1381.3687  ],\n",
       "          [ -780.93207 ,  -723.8065  ,  3132.8904  ],\n",
       "          ...,\n",
       "          [-1206.6353  ,   -81.008194,   -50.554405],\n",
       "          [  223.30663 ,   804.5833  ,   910.34674 ],\n",
       "          [  -59.93164 ,  -782.527   ,   679.01166 ]],\n",
       "  \n",
       "         [[-3503.1208  , -1257.9457  ,  2336.0732  ],\n",
       "          [ 1691.9426  ,  1231.0265  , -1076.2365  ],\n",
       "          [ -191.103   ,  1286.1288  , -1039.2635  ],\n",
       "          ...,\n",
       "          [  431.77588 ,  -657.54    ,   230.38184 ],\n",
       "          [ -647.2219  ,  -453.18    ,  -193.23723 ],\n",
       "          [  517.93945 ,  -114.20388 ,   330.99396 ]],\n",
       "  \n",
       "         [[-2213.3738  ,   472.26947 ,  1414.675   ],\n",
       "          [ -479.08417 ,  1277.0883  , -1584.1694  ],\n",
       "          [-1139.1154  ,  2544.798   ,  -507.48438 ],\n",
       "          ...,\n",
       "          [  313.10767 ,   736.7962  ,  -161.88942 ],\n",
       "          [ -641.8863  ,  -195.63199 ,   207.51501 ],\n",
       "          [ 1327.7233  , -3003.0706  ,    32.54847 ]]], dtype=float32),\n",
       "  'scale': array([1.94881939]),\n",
       "  'shift': array([-45185.30855141]),\n",
       "  'scaleshift_mode': 'atomwise',\n",
       "  'use_atomwise_scaleshift': array([ True]),\n",
       "  'avg_force_dis': array([0.00119437])})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_md17.save_dataset('dataset_ethanol', 1024, 128, 1024, mode='atomwise')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 ('mindsponge')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2496ecc683137a232cae2452fbbdd53dab340598b6e499c8995be760f3a431b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
